# -*- coding: utf-8 -*-
"""
Created on Fri Apr 14 12:19:38 2023

@author: lenovo
"""

import xarray as xr
import os
import numpy as np
import math
import scipy.interpolate
from osgeo import gdal
import re
import glob
# import shutil
from tqdm import tqdm # 运行计时

class para:
    climate_path=r'E:\drivedata\era5data'
    dem_path=r'E:\GDEM V3'
    slop_path=r'E:\slop30m'
    ndvi_path=r'E:\ndvi'
    landtypedir = r'F:\公司项目\陕西\商洛\land30m'
    out_path=r'F:\公司项目\陕西\商洛\data\GEP'
    inputyear=2023
    lon1=-1
    lon2=-1
    lat1=-1
    lat2=-1
    imlat=0
    imlon=0
    imwidth=3601
    imheight=3601
    num='0'
    soilr=1.274 #土壤容重，t/m3
################################################
def read_img(filename):
    dataset = gdal.Open(filename)  # 打开文件
    if dataset == None:
        print(filename+"文件无法打开")
        return
    im_width = dataset.RasterXSize  # 栅格矩阵的列数
    im_height = dataset.RasterYSize  # 栅格矩阵的行数
    im_bands = dataset.RasterCount #波段数
    im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
    im_proj = dataset.GetProjection()  # 地图投影信息
    im_data = dataset.ReadAsArray(0, 0, im_width, im_height).astype(np.float32)  # 将数据写成数组，对应栅格矩阵
    im_lon=[im_geotrans[0]+i*im_geotrans[1] for i in range(im_width)]
    im_lat=[im_geotrans[3]+i*im_geotrans[5] for i in range(im_height)]
    
    del dataset  # 关闭对象，文件dataset   
    return im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat
#=========================
# 保存tif文件函数
def write_img(im_data, im_geotrans, im_proj, path, nodata=None):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    # 创建文件
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
    if (dataset != None):
        dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
        dataset.SetProjection(im_proj)  # 写入投影        
    for i in range(im_bands):
        if (nodata != None):
            dataset.GetRasterBand(i + 1).SetNoDataValue(nodata)
        dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset

#=======================   
def readasc(fname):
    res=np.loadtxt(fname,skiprows=6)
    res[res==-9999]=np.nan
    return res
#======================================
def interpolation(lat,lon,data,met):
                
    latc=para.imlat
    lonc=para.imlon
    
    lon_c,lat_c=np.meshgrid(lonc,latc)

    lon_b,lat_b=np.meshgrid(lon,lat)
    #method{‘linear’, ‘nearest’, ‘cubic’}, optional
    res = scipy.interpolate.griddata((lon_b.ravel(), lat_b.ravel()),
                                    data.ravel(), (lon_c, lat_c), 
                                    method=met,fill_value=np.nan)
    return res
#==========================


# def get_C2(im_data):
#     im_height=np.size(im_data,0)
#     im_width=np.size(im_data,1)
#     cv=np.zeros((im_height,im_width))
#     for j in range(0,im_height):
#         for i in range(0,im_width):
#             vcf=(im_data[j,i]-0.02)/(1.0-0.02)*100
#             if vcf<=0:
#                 cv[j,i]=1
#             elif vcf<=78.3:
#                 cv[j,i]=0.6508-0.3436*(math.log10(vcf))
#             else:
#                 cv[j,i]=0
                
#             if cv[j,i]>1:
#                 cv[j,i]=1
#     return cv

def get_C(ndvi):   
    vcf=(ndvi-0.02)/(1.0-0.02)*100
    if vcf<=0:
        cv=1
    elif vcf<=78.3:
        cv=0.6508-0.3436*(math.log10(vcf))
    else:
        cv=0
        
    if cv>1:
        cv=1
            
    return cv


def get_K(sand,silt,clay,cc):
    san=sand
    sil=silt
    cla=clay
    sni=1-san/100
    
    a=0.2+0.3*math.exp(-0.0256*san*(1-sil/100))
    b=(sil/(cla+sil))**0.3
    # print(b,sil,cla)
        
    c=1-0.25*cc/(cc+math.exp(3.72-2.95*cc))
    d=1-0.7*sni/(sni+math.exp(-5.51+22.9*sni))
    
    kv=a*b*c*d*0.1317
    
    return kv

def get_LS(slop):
    sp=slop
    
    if sp<=1.0:
        m=2
    elif sp>1.0 and sp<=3.0:
        m=0.3
    elif sp>3.0 and sp<=5.0:
        m=0.4
    else:
        m=0.5
     
    ##方法2
    # sn=math.sin(sp/180*math.pi)
    # b=sn/0.089/(3*sn**0.8+0.56)
    # m=b/(1+b)
        
    length=30.0/math.cos(sp/180*math.pi)
    lv=(length/22.13)**m
    
    if sp<5.0:
        sv=10.8*math.sin(sp/180*math.pi)+0.03
    elif sp>=5.0 and sp<=10.0:
        sv=16.8*math.sin(sp/180*math.pi)-0.5
    else:
        sv=21.9*math.sin(sp/180*math.pi)-0.96
        
    return lv,sv

def get_R(pm):
    py=np.sum(pm)
    
    rv=0
    for p in pm:
        a=p**2
        b=1.5*math.log10(a/py)-0.8188
        c=1.735*10**b
        rv=rv+c
    rv=rv*17.02    
    return rv
    
def get_P(tp):
    if tp>=101 and tp<=108: #自然林地灌丛
        P=1
    elif tp>=109 and tp<=112: #人工乔木园地
        P=0.6
    elif tp>=21 and tp<=24: #草地
        P=0.9 
    elif tp>=31 and tp<=37: #河流湿地
        P=0
    elif tp==41: #水田
        P=0.15
    elif tp==42: #旱田
        P=0.35        
    elif tp>=51 and tp<=53: #9城镇
        P=0
    elif tp==63: #稀疏草地
        P=0.8
    else:
        P=1
        
    return P

#======================================
def get_soil(fn,k):
    
    # fn='H:\\全国\\1km\\T_clay_china.tif'
    im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat=read_img(fn)


    lon1=para.lon1
    lon2=para.lon2
    lat1=para.lat1
    lat2=para.lat2
    n1=-1
    n2=-1
    for n in range(0,im_width):
        if lon1<im_lon[n] and n1<0:
            n1=n-1
        if lon2<im_lon[n] and n2<0:
            n2=n
    if n2<0:
        n2=im_width-1

    t1=-1
    t2=-1
    for n in range(0,im_height):
        if lat1>im_lat[n] and t1<0:
            t1=n-1
        if lat2>im_lat[n] and t2<0:
            t2=n
    if t2<0:
        t2=im_height-1

    lat=im_lat[t1:t2+1]
    lon=im_lon[n1:n2+1]
    data=im_data[t1:t2+1,n1:n2+1]
    
    #利用全国平均数据补充空白区域
    if k==1:
        data[data<=0]=20.0 #clay
    else:
        data[data<=0]=30.0 #silt
    
    
    met='linear'
    res = interpolation(lat,lon,data,met)
    
    return res

def clc_R():
    monthday=np.array([31,28,31,30,31,30,31,31,30,31,30,31])
    path=para.climate_path
    # 提取指定经纬度范围内的数据
    lon_range = slice(para.lon1, para.lon2)  # 经度范围
    lat_range = slice(para.lat1, para.lat2)  # 纬度范围

    fn=path+'/Total_precipitation/total_precipitation'+str(para.inputyear)+'.nc'
    da=xr.open_dataset(fn)
    data = da.sel(longitude=lon_range, latitude=lat_range)

    lat=data.latitude.values
    lon=data.longitude.values
    pre=data.tp.values*1000 #m -> mm
    nx=np.size(pre,2)
    ny=np.size(pre,1)
    nd=np.size(pre,0)
    nd=int(nd/24)
    pre.shape=[24,nd,ny,nx]
    dpre=np.nansum(pre,axis=0)
    ri=np.zeros((ny,nx))
    pm=np.zeros(12)
    for i in range(0,nx):
        for j in range(0,ny):
            dd=0
            for m in range(0,12):
                sm=0.0
                for d in range(0,monthday[m]): 
                    sm=sm+dpre[dd,j,i]
                    dd=dd+1
                pm[m]=sm    
            ri[j,i]=get_R(pm)
            
    #插值到30m网格
    met='cubic'
    rv = interpolation(lat,lon,ri,met)
    
    return rv
                    
def clc_K():
    fn='E:\\1km\\T_clay_china.tif'
    clay=get_soil(fn,1)
    fn='E:\\1km\\T_silt_china.tif'
    silt=get_soil(fn,2)
    
    om=2.5 #土壤有机碳含量
    rows=para.imheight
    cols=para.imwidth
    
    kv=np.zeros((rows,cols))
    for i in range(0,rows):
        for j in range(0,cols):
            si=silt[i,j]
            cl=clay[i,j]                
            sa=100-si-cl
            kv[i,j]=get_K(sa,si,cl,om)
    return kv

def clc_LS():
    fn=para.slop_path+'/slop_N'+para.num[0]+'E'+\
        para.num[1]+'.tif'
        
    im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat=read_img(fn)
    
    lv=np.zeros((im_height,im_width))
    sv=np.zeros((im_height,im_width))
    for i in range(0,im_height):
        for j in range(0,im_width):
            l,s=get_LS(im_data[i,j])
            lv[i,j]=l
            sv[i,j]=s
    return lv,sv

def get_ndvi(fn):
    
    im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat=read_img(fn)


    lon1=para.lon1
    lon2=para.lon2
    lat1=para.lat1
    lat2=para.lat2
    n1=-1
    n2=-1
    for n in range(0,im_width):
        if lon1<im_lon[n] and n1<0:
            n1=n-1
        if lon2<im_lon[n] and n2<0:
            n2=n
    if n2<0:
        n2=im_width-1

    t1=-1
    t2=-1
    for n in range(0,im_height):
        if lat1>im_lat[n] and t1<0:
            t1=n-1
        if lat2>im_lat[n] and t2<0:
            t2=n
    if t2<0:
        t2=im_height-1

    lat=im_lat[t1:t2+1]
    lon=im_lon[n1:n2+1]
    data=im_data[t1:t2+1,n1:n2+1]
    
    data[np.isnan(data)] = 0
    data=data/10000
    data[data<0.2] = 0
    
    
    met='linear'
    res = interpolation(lat,lon,data,met)
    
    return res

def clc_C():
    fn=para.ndvi_path+'\\MOD13A2_'+str(para.inputyear)+'.tif'
    
    ndvi = get_ndvi(fn)
    
    rows=para.imheight
    cols=para.imwidth
    
    cv=np.zeros((rows,cols))
    for i in range(0,rows):
        for j in range(0,cols):
            if np.isnan(ndvi[i,j]):
                cv[i,j]=np.nan
            else:
                cc=get_C(ndvi[i,j])
                cv[i,j]=cc
    return cv

################################################

tiffile_all = glob.glob(para.landtypedir+'\\'+'*.tif')
for fn in tqdm(tiffile_all):
    tif_file_folder,tif_file_name = os.path.split(fn)
    para.num=re.findall('\d+', tif_file_name)
    
    im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat=read_img(fn)

    para.imlat=im_lat
    para.imlon=im_lon
    para.imheight=im_height
    para.imwidth=im_width
    para.lat1=im_lat[0]+0.5
    para.lat2=im_lat[-1]-0.5
    para.lon1=im_lon[0]-0.5
    para.lon2=im_lon[-1]+0.5
    


    cols=im_width
    rows=im_height
    pv=np.zeros((rows,cols))
    for i in range(0,rows):
        for j in range(0,cols):         
            pv[i,j]=get_P(im_data[i,j])

    rv=clc_R()
    kv=clc_K()
    lv,sv=clc_LS()
    cv=clc_C()

#------------------
    # rv[np.isnan(rv)]=255
    # fout='H:/全国/value/result/R_valur_test.tif'
    # write_img(rv, im_geotrans, im_proj, fout,255)

    # kv[np.isnan(kv)]=255
    # fout='H:/全国/value/result/K_value_test.tif'
    # write_img(kv, im_geotrans, im_proj, fout,255)

    # lv[np.isnan(lv)]=255
    # fout='H:/全国/value/result/L_value_test.tif'
    # write_img(lv, im_geotrans, im_proj, fout,255)

    # sv[np.isnan(sv)]=255
    # fout='H:/全国/value/result/S_value_test.tif'
    # write_img(sv, im_geotrans, im_proj, fout,255)

    # cv[np.isnan(cv)]=255
    # fout='H:/全国/value/result/C_value_test.tif'
    # write_img(cv, im_geotrans, im_proj, fout,255)
#------------------

    Qsr=rv*kv*lv*sv*(1-cv*pv)/10000 #t/m2/a    

    # cdr=140 #单位清淤成本，元/m3.
    cdr=35 #单位清淤成本，元/m3.
    Vsd=0.24*Qsr/para.soilr*cdr

    VdpdN=Qsr*0.37/100*1750 #贵州氮污染含量0.37%，价格1750元
    VdpdP=Qsr*0.108/100*2800 #贵州磷污染含量0.108%，价格2800元

    Vsr=Vsd+VdpdN+VdpdP
        

    Vsr[np.isnan(Vsr)]=0
    Vsr[Vsr<0]=0
    Vsr[im_data == 255] = -9999
    
    tif_file_name=tif_file_name.split('_')[1]
    fout=para.out_path+'\\土壤保持_'+tif_file_name
    write_img(Vsr, im_geotrans, im_proj, fout,-9999)
    
    # Qsr[Qsr==np.nan]=255
    # fout=para.out_path+'\\功能量\\Qsr_'+tif_file_name
    # write_img(Qsr, im_geotrans, im_proj, fout,255)


